import ujson
from usr.configure import settings
from usr.libs.logging import getLogger


logger = getLogger(__name__)


class CozeWebsocketConnection(object):

    def __init__(self, debug=True):
        self.debug = debug
        self.__recv_thread = None

    def __str__(self):
        return "{}".format(type(self).__name__)

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, *args, **kwargs):
        return self.disconnect()

    @property
    def conn(self):
        __client__ = getattr(self, "__client__", None)
        if __client__ is None:
            raise RuntimeError("{} not connected".format(self))
        return __client__
    
    def disconnect(self):
        """disconnect websocket"""
        __client__ = getattr(self, "__client__", None)
        if __client__ is not None:
            __client__.close()
            delattr(self, "__client__")
        if self.__recv_thread is not None:
            self.__recv_thread.join()
            self.__recv_thread = None

    def connect(self):
        """connect websocket"""
        __client__ = ws.Client.connect(
            "{}?bot_id=".format(settings.COZE_CHAT_URL, settings.COZE_BOT_ID),
            headers={
                "Authorization": "Bearer {}".format(settings.ACCESS_TOKEN),
                "Content-Type": "application/json"
            },
            debug=self.debug
        )
        try:
            self.__recv_thread = Thread(target=self.__recv_thread_worker)
            self.__recv_thread.start(stack_size=128)
        except Exception as e:
            __client__.close()
            logger.error("{} connect failed, Exception details: {}".format(self, repr(e)))
        else:
            setattr(self, "__client__", __client__)
            return __client__

    def __recv_thread_worker(self):
        while True:
            try:
                raw = self.conn.recv(1024*32)
            except Exception as e:
                logger.info("{} recv thread break, Exception details: {}".format(self, repr(e)))
                break
            if raw is None or raw == "":
                logger.info("{} recv thread break, Exception details: read none bytes, websocket disconnect".format(self))
                break
            if len(raw) > 1024*4:
                continue
            try:
                self.event_cb(ujson.loads(raw))
            except Exception as e:
                print("handle event error: {}".format(repr(e)))

    def event_cb(event):
        pass

    def emit(self, event_type, data):
        """发布客户端event"""
        return self.conn.send(
            ujson.dumps(
                {
                    "id": "{}<{}>".format(event_type, utime.mktime(utime.localtime())),
                    "event_type": event_type, 
                    "data": data
                }
            )
        )

    def chat_update(self):
        return self.emit(
            "chat.update",
            {
                "chat_config": {
                    "meta_data": {},
                    "custom_variables": {}, // 智能体中定义的变量。在智能体prompt中设置变量{{key}}后,可以通过该参数传入变量值,同时支持Jinja2语法。详细说明可参考变量示例。变量名只支持英文字母和下划线。
                    "auto_save_history": True,
                    "conversation_id": "xxxx", // conversation_id
                    "user_id": "xxx",  // 标识当前与智能体的用户，由使用方自行定义、生成与维护。user_id 用于标识对话中的不同用户，不同的 user_id，其对话的上下文消息、数据库等对话记忆数据互相隔离。如果不需要用户数据隔离，可将此参数固定为一个任意字符串
                    "extra_params": {},   // 附加参数,通常用于特殊场景下指定一些必要参数供模型判断,例如指定经纬度,并询问智能体此位置的天气。自定义键值对格式,其中键(key)仅支持设置为:latitude:纬度,此时值(Value)为纬度值,例如39.9800718。longitude:经度,此时值(Value)为经度值,例如116.309314。
                    "parameters": {"custom_var_1": "测试"}
                },
                "input_audio": {         // 输入音频格式
                    "format": "pcm",       // 输入音频格式，支持 pcm/wav/ogg。默认 wav
                    "codec": "pcm",         // 输入音频编码。 pcm/opus。默认 pcm
                    "sample_rate": 24000,  // 采样率
                    "channel": 1, // 通道数
                    "bit_depth": 16 // 位深
                },
                "output_audio": {        // 输出音频格式
                    "codec": "pcm",       
                    "pcm_config": {
                        "sample_rate": 16000,  // 默认  24000
                        "frame_size_ms": 50,
                        "limit_config": {
                            "period": 1,
                            "max_frame_num": 22
                        }
                    },
                    "speech_rate": 0,  // 回复的语速，取值范围 [-50, 100]，默认为 0，-50 表示 0.5 倍速，100 表示 2倍速
                    "voice_id": "7426720361733046281"
                }
            }
        )
